import torch
from torch.nn import Module, Linear, ModuleList, LayerNorm, Sequential, GELU
from seasoncnnblock import CnnSeasonBlock

class CnnSeason(Module):
    def __init__(self, layer:int, topk:int, input_features:int, model_dim:int, seq:int, target_id:int):
        super(CnnSeason, self).__init__()
        self.topk = topk
        self.target_id = target_id
        self.embed = Linear(input_features, model_dim)
        self.layer = layer
        self.models = ModuleList([CnnSeasonBlock(self.topk, model_dim) for _ in range(layer)])
        self.layer_norm = LayerNorm(model_dim)
        self.dembed = Linear(model_dim, 1)
        self.perdict = Sequential(Linear(seq, seq), GELU(), Linear(seq, 1))


    def forward(self, season:torch.Tensor):
        # 先norm
        means = season.mean(dim=1, keepdim=True).detach()
        season -= means
        stdev = torch.sqrt(torch.var(season, dim=1, keepdim=True, unbiased=False) + 1e-5)
        season /= stdev

        # 对season进行特征embedding，将特征数从input_features变为model_dim
        embeded = self.embed(season)

        # cnnseasonblock层
        for i in range(self.layer):
            embeded = self.layer_norm(self.models[i](embeded))
        
        # 将特征集中到一起，只留要预测的目标特征
        output = self.dembed(embeded).squeeze(-1)
        
        # denorm，用目标特征的均值方差来做
        output = output * stdev[:,:,self.target_id] + means[:,:,self.target_id]
        
        # 将时间特征集中在一起，得到目标时间点
        output = self.perdict(output).squeeze(-1)
        return output


if __name__ == "__main__":
    B = 2
    T = 13
    N = 4
    x = torch.ones(B, T, N)
    block = CnnSeason(layer=3, topk=3, input_features=4, model_dim=12, seq=T, target_id=0)
    res = block(x)